package online.heycm.platform.web.common.filter.wrapper;

import org.springframework.util.CollectionUtils;
import org.springframework.util.FastByteArrayOutputStream;
import org.springframework.util.StringUtils;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.StringJoiner;

/**
 * 缓存响应体
 *
 * @author hey
 * @version 1.0
 * @date 2024/1/5 15:57
 */
public class CacheResponseWrapper extends HttpServletResponseWrapper {

    private FastByteArrayOutputStream fastByteArrayOutputStream;
    private ServletOutputStream servletOutputStream;
    private PrintWriter printWriter;
    // 响应头
    private Map<String, String> extHeaders;
    // 状态码
    private int status = -1;

    /**
     * Constructs a response adaptor wrapping the given response.
     *
     * @param response the {@link HttpServletResponse} to be wrapped.
     * @throws IllegalArgumentException if the response is null
     */
    public CacheResponseWrapper(HttpServletResponse response) {
        super(response);
        fastByteArrayOutputStream = new FastByteArrayOutputStream(1024);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        if (servletOutputStream == null) {
            servletOutputStream = new CacheServletOutputStream(fastByteArrayOutputStream);
        }
        return servletOutputStream;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (printWriter == null) {
            String encoding = getCharacterEncoding();
            if (StringUtils.isEmpty(encoding)) {
                encoding = "UTF-8";
            }
            ServletOutputStream outputStream = getOutputStream();
            printWriter = new PrintWriter(new OutputStreamWriter(outputStream, encoding));
        }
        return printWriter;
    }

    /**
     * 获取已写入数据的字节数组
     */
    public byte[] getContentAsByteArray() {
        return fastByteArrayOutputStream.toByteArray();
    }

    /**
     * 获取已写入数据的字符
     */
    public String getContentAsString() {
        return new String(getContentAsByteArray(), StandardCharsets.UTF_8);
    }

    /**
     * 将缓冲区内容写回到原始ServletOutputStream并发送给客户端
     */
    public void writeToClient() throws IOException {
        // 缓冲区数据
        byte[] content = getContentAsByteArray();
        // 重置原始HTTP响应
        HttpServletResponse rawResponse = (HttpServletResponse) getResponse();
        rawResponse.reset();
        String contentType = getContentType();
        if (contentType != null) {
            rawResponse.setContentType(contentType);
        }
        rawResponse.setContentLength(content.length);
        for (Map.Entry<String, String> entry : getExtHeaders().entrySet()) {
            String headerName = entry.getKey();
            String headerValue = entry.getValue();
            if (headerValue.contains(",")) {
                String[] values = headerValue.split(",");
                for (String v : values) {
                    rawResponse.addHeader(headerName, v.trim());
                }
            } else {
                rawResponse.setHeader(headerName, headerValue);
            }
        }
        if (this.status != -1) {
            rawResponse.setStatus(status);
        }
        ServletOutputStream outputStream = rawResponse.getOutputStream();
        outputStream.write(content);
        outputStream.flush();
    }

    @Override
    public void setHeader(String name, String value) {
        super.setHeader(name, value);
        addExtHeader(name, value);
    }

    @Override
    public void setIntHeader(String name, int value) {
        super.setIntHeader(name, value);
        addExtHeader(name, value + "");
    }

    @Override
    public void setDateHeader(String name, long date) {
        super.setDateHeader(name, date);
        addExtHeader(name, date + "");
    }

    @Override
    public void addHeader(String name, String value) {
        super.addHeader(name, value);
        addExtHeader(name, getRawHeaderValues(name));
    }

    @Override
    public void addDateHeader(String name, long date) {
        super.addDateHeader(name, date);
        addExtHeader(name, getRawHeaderValues(name));
    }

    @Override
    public void addIntHeader(String name, int value) {
        super.addIntHeader(name, value);
        addExtHeader(name, getRawHeaderValues(name));
    }

    @Override
    public void setStatus(int sc) {
        super.setStatus(sc);
        this.status = sc;
    }

    @Override
    public void setStatus(int sc, String sm) {
        super.setStatus(sc, sm);
        this.status = sc;
    }

    public Map<String, String> getExtHeaders() {
        return extHeaders == null ? Collections.emptyMap() : extHeaders;
    }

    private void addExtHeader(String name, String value) {
        if (extHeaders == null) {
            extHeaders = new HashMap<>();
        }
        extHeaders.put(name, value);
    }

    private void addExtHeader(String name, Collection<String> values) {
        if (values != null && !values.isEmpty()) {
            StringBuilder sb = new StringBuilder();
            for (String v : values) {
                sb.append(v).append(",");
            }
            addExtHeader(name, sb.toString().trim());
        }
    }

    private Collection<String> getRawHeaderValues(String name) {
        return ((HttpServletResponse) getResponse()).getHeaders(name);
    }

    private class CacheServletOutputStream extends ServletOutputStream {

        private final OutputStream os;

        private CacheServletOutputStream(OutputStream os) {
            this.os = os;
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }

        @Override
        public void write(int b) throws IOException {
            this.os.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            this.os.write(b, off, len);
        }

        @Override
        public void flush() throws IOException {
            this.os.flush();
        }

        @Override
        public void close() throws IOException {
            this.os.close();
        }
    }
}
